#coding:utf-8
import paramiko
import sqlite3
import traceback
import base64
from multiprocessing.dummy import Pool as ThreadPool
import os
import demjson
import requests
ConsoleIP = ''

def setIP():
    global ConsoleIP
    try:
        ConsoleIP = requests.get('http://ip.42.pl/raw').text
    except Exception as e:
        print()
        print('获取IP异常')
        raise
setIP()
BASEURL = 'http://'+ConsoleIP+':5000'
print BASEURL

#执行单句有返回
def DB_EXECSQL_ONE(SQL):
    try:
        conn = sqlite3.connect('console.db')
        c = conn.cursor()
        c.execute(SQL)
        res = c.fetchall()
        conn.commit()
        conn.close()
        return res
    except Exception as e:
        print('DB_EXECSQL_ONE:',SQL,e)
        return False
#批量执行SQL
def DB_EXECSQL(SQLLIST):
    try:
        conn = sqlite3.connect('console.db')
        c = conn.cursor()
        for SQL in SQLLIST:
            c.execute(SQL)
            conn.commit()
        conn.close()
        return True
    except Exception as e:
        print('DB_EXECSQL:',SQLLIST,e)
        return False
#数据库初始化
def INIT():
    #数据库初始化
    if os.path.exists('console.db') == False:
        CreateServersTableSQL = '''CREATE TABLE SERVERS (
            ID       INTEGER       PRIMARY KEY AUTOINCREMENT
                                   NOT NULL,
            IP       VARCHAR (20)  UNIQUE
                                   NOT NULL,
            PASSWORD VARCHAR (100) NOT NULL,
            STATUS   INTEGER       NOT NULL
        );
        '''
        # 创建任务表
        CreateTaskTableSQL = '''CREATE TABLE TASK (
            ID      INTEGER PRIMARY KEY AUTOINCREMENT
                            NOT NULL,
            NAME    VARCHAR NOT NULL,
            CONFIG  TEXT    NOT NULL,
            RESULTS TEXT
        );
        '''
        # 创建任务回调日志表
        CreateCallBackLogsTableSQL = '''CREATE TABLE TASKLOGS (
            ID     INTEGER      PRIMARY KEY AUTOINCREMENT,
            TASKID INTEGER,
            IP     VARCHAR (20),
            LOG    TEXT
        );
        '''
        SQLLIST = []
        SQLLIST.append(CreateServersTableSQL)
        SQLLIST.append(CreateTaskTableSQL)
        SQLLIST.append(CreateCallBackLogsTableSQL)
        DB_EXECSQL(SQLLIST)
    #项目文件夹初始化
    if os.path.exists('privatekeys') == False:
        os.mkdir('privatekeys')
    if os.path.exists('static/ipfiles') == False:
        os.makedirs('static/ipfiles')

INIT() #项目初始化
print("Project Init......")

def shell_exec(commands):
    try:
        res = os.popen(commands)
        return res.read()
    except Exception as e:
        print(e)
        return False
def wc(filename):
    try:
        res = shell_exec('wc -l '+filename)
        return  int(res.split(' ')[0])
    except Exception as e:
        print(e)
        return False
#执行命令
def SSH_Exec_Command(IP,PASSWORD,COMMAND):
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if os.path.exists('privatekeys/'+PASSWORD) == True:
            private_key = paramiko.RSAKey.from_private_key_file('privatekeys/'+PASSWORD,password='admins')
            ssh.connect(hostname=IP, port=22, username='root', pkey=private_key)
        else:
            ssh.connect(hostname=IP, port=22, username='root', password=PASSWORD)
        stdin, stdout, stderr = ssh.exec_command(COMMAND)
        res, err = stdout.read(), stderr.read()
        result = res if res else err
        ssh.close()
        return result.decode()
    except Exception as e:
        traceback.print_exc()
        print('SSH_Exec_Command:',IP,PASSWORD,COMMAND,e)
        return False
#--------------扫描机管理--------------------
#增加服务器
def addServer(IP,PASSWORD):
    try:
        RETURNSTR = str(SSH_Exec_Command(IP,PASSWORD,'whoami'))
        if RETURNSTR.find('root') > -1:
            #密码base64后再入库,本段代码不兼容PY3
            SQL = "INSERT INTO SERVERS VALUES (null,'%s','%s',0);" % (IP,base64.b64encode(PASSWORD))
            DB_EXECSQL([SQL])
            systemInit(IP,PASSWORD) #验证并入库后立刻部署环境
            saveLogs('addServer success:' + IP + '  ' + PASSWORD)
            return True
        else:
            print(IP+'ssh connect failed......')
            saveLogs('addServer error:' + IP + '  ' + PASSWORD + '  ')
            return False
    except Exception as e:
        traceback.print_exc()
        saveLogs('Exception addServer error:'+IP+'  '+PASSWORD)
        print('addServer:',IP,PASSWORD,e)
        return False

def getpwd():
    return shell_exec('pwd').replace('\n','')
#删除服务器
def delServer(IP):
    try:
        SQL = "DELETE FROM SERVERS WHERE IP = '%s';" % IP
        DB_EXECSQL([SQL])
        return True
    except Exception as e:
        print('delServer:',IP,e)
        return False
#查询服务器
def getServers():
    SQL = "SELECT * FROM SERVERS;"
    serverList = DB_EXECSQL_ONE(SQL)
    return serverList
#环境初始化
def systemInit(IP,PASSWORD):
    #客户端装好以后 环境下会有一个install.lock内容是ok
    #检查是否存在环境
    RETURNSTR = str(SSH_Exec_Command(IP,PASSWORD,'''cat /var/client/install.lock'''))
    if RETURNSTR.find('ok') > -1:
        saveLogs(IP+"env found")
        print(IP,"env found")
        return True
    else:
        print(IP,"env not found,do install.")
        saveLogs(IP + "env not found,do install.")
        temp = '''ps -ef | grep "python" | grep -v grep | awk '{print $2}' | xargs kill -9\n'''
        temp += '''ps -ef | grep "zmap" | grep -v grep | awk '{print $2}' | xargs kill -9\n'''
        temp += '''ps -ef | grep "zgrab" | grep -v grep | awk '{print $2}' | xargs kill -9\n'''
        temp += 'cd /var\n'
        temp += 'rm -rf /var/client.zip\n'
        temp += 'apt-get -y install wget unzip zip\n'
        temp += 'wget -P /var http://'+ConsoleIP+'/client.zip\n'
        temp += 'unzip /var/client.zip\n'
        temp += 'rm -rf /var/client.zip\n'
        temp += 'cd /var/client\n'
        temp += 'wget http://'+ConsoleIP+'/install.sh\n'
        temp += 'chmod 777 -R *\n'
        temp += 'nohup sh /var/client/install.sh >> /var/client/install.log 2>&1 &\n\n'
        command = 'echo -e "%s" > /root/i.sh;' % temp
        command += 'nohup sh /root/i.sh > /root/i.log 2>&1 &'
        SSH_Exec_Command(IP, PASSWORD, command)
        return False

def saveLogs(log):
    logfile = open('static/log.txt','a+')
    logfile.write(log+'\n')
    logfile.close()
#--------------任务管理----------------------
def Task(info):
    try:
        url = info['spider']
        payloads = {
            'taskInfo':info['CONFIG']
        }
        response = requests.post(url,data=payloads)
        res = demjson.decode(response.text)
        print('TASK:',url,response.text)
        if res['status'] == 'success':
            saveLogs(info['spider'] + '   ->   success')
        else:
            saveLogs(info['spider'] + '   ->   failed')
    except Exception as e:
        traceback.print_exc()
        saveLogs(info['spider'] + '   ->   failed'+'\n'+e)
        print(e)
#下达任务
def addTask(NAME,CONFIG):
    try:
        serverList = getServers()
        serverCount = len(serverList)
        spiders = []
        for server in serverList:
            IP = server[1]
            spiders.append("http://" + IP + ':5000/spider')

        pwd = getpwd()
        SplitIPDIR = pwd + '/static/ipfiles/'
        CONF = demjson.decode(CONFIG)
        if CONF['SSL'] == 'False':
            IPFILE = pwd + '/static/80.txt'
        else:
            IPFILE = pwd + '/static/443.txt'
        print(IPFILE)
        if os.path.exists(IPFILE) == False:
            print('Can\'t Found IP FILE!')
            return False

        os.system('rm -rf '+ SplitIPDIR)
        os.system('mkdir ' + SplitIPDIR)

        IPCount = wc(IPFILE)
        avg = int(IPCount / serverCount) + 10

        shell_exec('split -l %s %s %ssplitip_' % (avg,IPFILE,SplitIPDIR)) #执行分割

        SplitIPFiles = os.listdir(SplitIPDIR)

        ipurls = []

        SQL = "INSERT INTO TASK VALUES (null,'%s','%s','[]');" % (NAME, base64.b64encode(CONFIG))
        print(SQL)
        DB_EXECSQL([SQL])

        SQL = "SELECT MAX(ID) FROM TASK;"
        res = DB_EXECSQL_ONE(SQL)
        TaskID = str(res[0][0])

        TASKS = []

        for ipfile in SplitIPFiles:
            ipurls.append(BASEURL+'/static/ipfiles/'+ipfile)
        for i,ipurl in enumerate(ipurls):
            #组合任务队列
            CONF['TASKID'.decode()] = TaskID.decode()
            CONF['IPFileUrl'.decode()] = ipurl.decode()
            CONF['MonitorCallBackUrl'.decode()] = str(BASEURL+'/addTaskLogs').decode()
            CONF['ResultUrl'.decode()] = str(BASEURL+'/addResult').decode()

            print(CONF)

            task = {
                'spider':spiders[i],
                'CONFIG':demjson.encode(CONF)
            }

            TASKS.append(task)
        #并发下发任务
        pool = ThreadPool(5)
        pool.map(Task, TASKS)
        pool.close()
        pool.join()
        return True
    except Exception as e:
        traceback.print_exc()
        print(e)
        return False
def submitZgrab(task):
    try:
        res = requests.post(task['zgraber'],data={'config':task['info']})
        print task['zgraber'],'---->',res.text
    except Exception as e:
        traceback.print_exc()
        print e
def Zgrab(config):
    temp = demjson.decode(config)
    pwd = getpwd()
    SplitIPDIR = pwd + '/static/ipfiles/'
    os.system('rm -rf ' + SplitIPDIR)
    os.system('mkdir ' + SplitIPDIR)
    IPFILE = pwd + '/static/'+temp['Port']+'.txt'
    if os.path.exists(IPFILE) == False:
        return 'IP FILE ERROR'
    else:
        serverList = getServers()
        serverCount = len(serverList)
        zgrabs = []
        for server in serverList:
            IP = server[1]
            zgrabs.append("http://" + IP + ':5000/zgrab')
        IPCount = wc(IPFILE)
        avg = int(IPCount / serverCount) + 10
        shell_exec('split -l %s %s %ssplitip_' % (avg, IPFILE, SplitIPDIR))  # 执行分割
        SplitIPFiles = os.listdir(SplitIPDIR)
        TASKS = []
        for i,ipfile in enumerate(SplitIPFiles):
            taskinfo = demjson.encode({
                'ContainList': temp['ContainList'],
                'ExcludeList': temp['ExcludeList'],
                'IPFile': BASEURL+'/static/ipfiles/'+ipfile,
                'Port': temp['Port'],
                'ReqFile': temp['ReqFile'],
                'Thread': temp['Thread'],
                'TimeOut': temp['TimeOut'],
                'TLS': temp['TLS']
            })
            zgraber = zgrabs[i]
            task = {
                'info':taskinfo,
                'zgraber':zgraber
            }
            TASKS.append(task)
        pool = ThreadPool(5)
        pool.map(submitZgrab, TASKS)
        pool.close()
        pool.join()
#获取任务列表
def getTaskList():
    SQL = "SELECT * FROM TASK;"
    TaskList = DB_EXECSQL_ONE(SQL)
    return TaskList
#任务状态回调
def getTaskLogs(TaskID):
    #根据任务ID获取回调日志
    SQL = "SELECT * FROM TASKLOGS WHERE TASKID = %s;" % TaskID
    TaskLogs = DB_EXECSQL_ONE(SQL)
    return TaskLogs

def addTaskLogs(TaskID,IP,Log):
    #进行BASE64编码后入库
    try:
        SQL = "DELETE FROM TASKLOGS WHERE TASKID = %s AND IP = '%s';" % (TaskID,IP)
        DB_EXECSQL([SQL])
        SQL = "INSERT INTO TASKLOGS VALUES (null,'%s','%s','%s');" % (TaskID,IP, base64.b64encode(Log))
        DB_EXECSQL([SQL])
        return True
    except Exception as e:
        print(e)
        return False

def addResult(TaskID,Result):
    try:
        SQL = "SELECT RESULTS FROM TASK WHERE ID = %s;" % TaskID
        res = DB_EXECSQL_ONE(SQL)
        res = demjson.decode(res[0][0])
        res.append(Result)
        res = demjson.encode(res)
        SQL =  "UPDATE TASK SET RESULTS = '%s' WHERE ID = %s;" % (res,TaskID)
        DB_EXECSQL([SQL])
        return True
    except Exception as e:
        print(e)
        return  False

def MakeCIDRIP():
    cidrip = []
    for i in range(1,225):
        if i not in [10,127]:
            cidrip.append(str(i)+'.0.0.0/8')
    return cidrip
def splitList(l,c):
    rl = []
    r = len(l) % c #余数
    n = l[0:len(l) - r]
    x = l[len(n):len(l)]
    avg = int(len(n) / c)
    for i in range(0,c):
        d = l[avg*i:avg*(i+1)]
        rl.append(d)
    for i,ii in enumerate(x):
        rl[i].append(ii)
    return rl
def dozmap(task):
    try:
        res = requests.post('http://'+task['ip']+':5000/zmap',data={
            'cidr':task['cidr'],
            'port':task['port'],
            'Band':task['Band']
        })
        print(task,res.text)
    except Exception as e:
        print(e,task)
def batchZmap(port,Band):
    try:
        servers = getServers()
        count = len(servers)
        cidrs = MakeCIDRIP()
        ips = splitList(cidrs,count)
        tasklist = []
        for i,ip in enumerate(ips):
            task = {
                'ip':servers[i][1],
                'cidr':demjson.encode(ip),
                'port':port,
                'Band':Band
            }
            tasklist.append(task)
        pool = ThreadPool(5)
        pool.map(dozmap, tasklist)
        pool.close()
        pool.join()
    except Exception as e:
        traceback.print_exc()
        print(e)
def bshell(obj):
    command = 'echo -e "%s" > /var/www/html/bshell.sh;' % obj['commands']
    command += 'nohup sh /var/www/html/bshell.sh > /var/www/html/bshell.txt 2>&1 &'
    SSH_Exec_Command(obj['ip'], obj['password'], command)
def batchShell(shellcode):
    try:
        #批量执行shell
        servers = getServers()
        objs = []
        for server in servers:
            SSH_Exec_Command(server[1], base64.b64decode(server[2]), shellcode)
        return True
    except Exception as e:
        print('batchShell',e)
        return False